import os
import argparse

import numpy as np
import h5py


def flip_dataset(hdf5_file, flip_file):
    vol_dataset = h5py.File(hdf5_file)
    flip_list = np.load(flip_file)
    print(len(flip_list))

    for name in flip_list:
        if not vol_dataset.__contains__(name):
            continue
        vol = vol_dataset[name][()]
        attrs = dict(vol_dataset[name].attrs)
        vol = vol[:, :, ::-1]
        del vol_dataset[name]
        vol_data = vol_dataset.create_dataset(name, data=vol, dtype=np.int16,
                                              chunks=(512, 512, 1), compression='gzip')
        for key in attrs:
            vol_data.attrs[key] = attrs[key]

        print('%s is reversed.' % name)

    vol_dataset.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Flip dataset according to flip_list")
    parser.add_argument("-d", "--dataset",
                        dest="hdf5_file", metavar="hdf5 file path",
                        help="target dataset to flip")
    parser.add_argument("-l", "--list",
                        dest="flip_file", metavar="flip list file path",
                        default='./flip_list.npy', required=False,
                        help="target scans to flip")

    args = parser.parse_args()

    flip_dataset(args.hdf5_file, args.flip_file)

